use ndarray::{Array1, Array2};

pub fn solve_by_gcr(a: Array2<f64>, b: Array1<f64>, tol: f64, max_iter: usize) -> Array1<f64> {
    let n = a.nrows();
    assert_eq!(a.ncols(), n, "Matrix A must be square.");
    assert_eq!(b.len(), n, "Vector b size must match matrix dimensions.");

    let mut x = Array1::<f64>::zeros(n); // Initial guess x = 0
    let mut r = b - &a.dot(&x);          // Residual r = b - Ax
    let mut p = r.clone();              // Search direction p = r
    let mut res_norm = r.dot(&r).sqrt(); // Initial residual norm

    for _ in 0..max_iter {
        // Stop if the residual is below the tolerance
        if res_norm < tol {
            break;
        }

        // Compute alpha = (rᵀ * r) / (pᵀ * A * p)
        let ap = a.dot(&p);
        let alpha = r.dot(&r) / p.dot(&ap);

        // Update the solution: x = x + alpha * p
        x = x + &(&p * alpha);

        // Update the residual: r = r - alpha * A * p
        let r_next = &r - &(ap * alpha);

        // Check for convergence
        res_norm = r_next.dot(&r_next).sqrt();
        if res_norm < tol {
            break;
        }

        // Compute beta = (r_nextᵀ * r_next) / (rᵀ * r)
        let beta = r_next.dot(&r_next) / r.dot(&r);

        // Update search direction: p = r_next + beta * p
        p = &r_next + &(p * beta);

        // Move to the next residual
        r = r_next;
    }

    x
}

#[cfg(test)]
#[allow(unused)]
mod test {
    use super::*;
    use ndarray::array;
    use approx::assert_abs_diff_eq;

    #[test]
    fn test_solve() {
        let a = array![
            [4.0, 1.0, 2.0],
            [3.0, 5.0, 1.0],
            [1.0, 1.0, 3.0]
        ];
        let b = array![4.0, 7.0, 3.0];

        let x = solve_by_gcr(a, b, 1e-6, 1000);
        
        let ans = array![0.5, 1.0, 0.5];
        
        assert_abs_diff_eq!(
            x.as_slice().unwrap(), 
            ans.as_slice().unwrap(),
            epsilon = 1e-9
        );
    }
}